import random,threading,ga_agent
import numpy,numpy.linalg,copy, time,pdb,pickle
class ga_agent_factory(object):
    def __init__(self,game_model,state_in_hidden_out,team):
        self.gm=game_model
        self.team=team
        # Maintenance settings and variables
        self.genetic_maintenance_time=20 # Time (sec.) between kill/breed cycles
        self.brain_lock=threading.Lock()
        self.maintenance_thread=threading.Thread(target=self.__maintenance_thread)
        self.maintenance_thread_enabled=True
        # Neural Network Nettings
        self.state_in_hidden_out=state_in_hidden_out # For a given state we have a list of the number of input, hidden, and output nodes for them
        self.initial_weight_std=1
        self.abs_max_weight=float('1e100')
        # Genetic Algorithm Settings
        self.num_brains=10
        self.select_fit_prob=.7
        self.kill_weak_prob=.9
        self.per_state_crossover_prob=.3
        self.per_weight_mutation_prob=.2
        self.mutation_weight_std=.2
        self.new_member_prob=.1
        self.min_mature_agents=round(self.num_brains*.25) # Minimum mature agents required to allow breeding to occur
        self.max_brains_removed_per_cycle=round(self.num_brains*.5)
        self.num_elite_agents=round(self.num_brains*.2)
        #self.max_family_members=round(self.num_brains*.2)
        self.min_mind_swap_time=5.0#10 # Max speed (sec.) that minds can be swapped
        self.min_mature_time=self.min_mind_swap_time*2.0#20 # Minimum time a brain has to be in a body before it can be evaluated
        self.fitness_learning_rate=.3 # Exponential time constant for averaging fitness values
        # Genetic Algorithm variables
        self.cur_family_id=0 # Each time a new brain is created, it is assigned to this value which is then incremented.
        self.cur_brain_index=0 # This is the next brain to be used in a body
        # Genetic Algorithm Structure (all are one per brain, the brain's position will be consistent for the lifetime of the brain)
        self.brains=[None]*self.num_brains # The [state][to_layer][to_node][from_node and bias] structure for the neural network weights (multilayer, feedforward)
        self.seconds_used=[None]*self.num_brains # The total time a brain has been in a body
        self.family=[None]*self.num_brains # Family ID number generated when a new brain is created (during crossover it takes the "dominant" parents)
        self.fitness=[None]*self.num_brains # This is the average fitness for each brain for each fitness metric (which there may be many of, to get the scalar fitness, we request the current fitness weights from the game model and it will give us a vector that we can take the dot product with this and get our true fitness)
        self.true_fitness=[float('-1e1000')]*self.num_brains # This is the scalar valued fitness, it must be updated every maintenance period
        self.checked_out=[None]*self.num_brains # True if there is an  outstanding brain 'checkout', we don't allow brains to be killed that are checked out
        self.death_mark=[None]*self.num_brains # When this is true for a brain, we will remove it during maintenance whenever it has no brains_checked_out.  This means no fitness information is lost.  When true, new brains can't be checked out.
        
        for i in range(self.num_brains):
            self.initialize_cur_family(i)
            self.initialize_brain_aux(i)
            self.brains[i]=self.__make_new_brain()
        self.create_maintenance_thread()

    def initialize_brain_aux(self,num):
        self.seconds_used[num]=0.0
        self.fitness[num]=[]
        self.true_fitness[num]=0
        self.checked_out[num]=False
        self.death_mark[num]=False
    def initialize_cur_family(self,num):
        self.family[num]=self.cur_family_id
        self.cur_family_id+=1

    ##########################################################
    # External Object Creation/Updating, saving/loading and Thread Control
    def save_state(self):
        self.brain_lock.acquire()
        file_name=str(time.time())+'-team'+str(self.team)+'.brains'
        pickle.dump((self.brains,self.seconds_used,self.family,self.fitness,self.true_fitness,self.checked_out,self.death_mark,self.num_brains),file(file_name,'w'))
        self.brain_lock.release()
    def load_state(self,file_name):
        self.brain_lock.acquire()
        try:
            self.brains,self.seconds_used,self.family,self.fitness,self.true_fitness,checked_out,self.death_mark,self.num_brains=pickle.load(file(file_name,'r'))
        except IOError,tmp:
            print('File not found when loading GA state: '+file_name)
        self.brain_lock.release()
    def create_agent_slot(self,actions,senses,object_type,object_num,max_velocity,max_heading_diff,radius,graphics_draw,graphics_die,pos,heading,im_file):
        """Done in the beginning to create initial agents, instantiates the ga_agent with a brain, the current time, and the required team information, sensors, etc"""
        # Create a GA agent with the required properties
        self.brain_lock.acquire()
        self.goto_valid_brain_index()
        self.checked_out[self.cur_brain_index]=True
        new_agent=ga_agent.ga_agent(self,body_time=time.time(),brain_index=self.cur_brain_index,brain=self.brains[self.cur_brain_index],state_in_hidden_out=self.state_in_hidden_out,actions=actions,senses=senses,object_type=object_type,object_num=object_num,max_velocity=max_velocity,max_heading_diff=max_heading_diff,radius=radius,graphics_draw=graphics_draw,graphics_die=graphics_die,pos=pos,heading=heading,im_file=im_file)

        self.inc_brain_index()
        self.brain_lock.release()
        return new_agent
    def compute_agents_statistics(self,agent):
        """Calls various updating functions so that the fitness vector is correct.  It is necessary as, though the game model calls refresh agent, other things can too (e.g., the agent itself), thus we need a way of getting all of our information at the exact moment we want to switch brains"""
        agent_num=agent.object_num
        self.gm.fit_sbm_update_fitness_agent(agent_num)
        temp_vec=agent.fitness_vector
        agent.clear_fitness_stats()
        self.gm.fit_sbm_new_holder_agent(agent_num)
        return temp_vec
    def refresh_agent(self,agent):
        """Called whenever the game model wants to change out the brain of an agent, this will call the fitness function for this agent, and update the overall fitness for the 'mind'.  If the brain hasn't been in the body for at least min_mind_swap_time, we return.  If we are busy doing maintenance, we return."""
        if not self.brain_lock.acquire(False):
            return
        if not agent.body_lock.acquire(False):
            self.brain_lock.release()
            return
        if time.time()-agent.body_time < self.min_mind_swap_time:
            self.brain_lock.release()
            agent.body_lock.release()
            return
        # Get the information we need from the current brain
        brain_num=agent.brain_index
        self.seconds_used[brain_num]+=time.time()-agent.body_time
        
        agent_fitness=copy.deepcopy(self.compute_agents_statistics(agent))
        old_fit=self.fitness[brain_num]
        if len(self.fitness[brain_num])==0:
            self.fitness[brain_num]=agent_fitness
        else: # Combine fitness vectors
            self.fitness[brain_num]=map(lambda old_fit, new_fit: old_fit*(1.0-self.fitness_learning_rate)+new_fit*self.fitness_learning_rate,self.fitness[brain_num],agent_fitness)
        print('Old Fit: %s New Fit: %s') %(str(old_fit),str(self.fitness[brain_num]))
        
        # Cleanup body for use by new brain
        self.checked_out[brain_num]=False

        # Pick a new brain
        self.goto_valid_brain_index()
        new_brain=self.cur_brain_index
        print('Swaping out brain %d for brain %d') %(brain_num,new_brain)
        self.checked_out[new_brain]=True
        agent.brain_index=new_brain
        agent.brain=self.brains[new_brain]
        agent.body_time=time.time()

        self.inc_brain_index()
        agent.body_lock.release()
        self.brain_lock.release()

    def create_maintenance_thread(self):
        """Creates a thread that performs the kill/breed cycle every genetic_maintenance_time seconds"""
        self.maintenance_thread_enabled=True
        self.maintenance_thread.start()

    def stop_maintenance_thread(self):
        """Kills the maintenance thread"""
        self.maintenance_thread_enabled=False
        self.maintenance_thread.join()

    ##########################################################
    # Internal Maintenance Functions
    def goto_valid_brain_index(self):
        while self.death_mark[self.cur_brain_index] or self.checked_out[self.cur_brain_index] > 0:
            self.inc_brain_index()

    def inc_brain_index(self):
        self.cur_brain_index+=1
        if self.cur_brain_index==self.num_brains:
            self.cur_brain_index=0
    
    def compute_fitness_list(self,agents):
        """Returns a scalar list of fitness values by taking each agents fitness vector and dotting it with the fitness weights (from the game model)"""
        self.gm.fit_weights_lock.acquire()
        fit_weights=numpy.array(self.gm.fit_weights[self.team])
        for agent in agents:
            self.true_fitness[agent]=numpy.dot(numpy.array(self.fitness[agent]),fit_weights)
        self.gm.fit_weights_lock.release()

    def __maintenance_thread(self):
        """This waits until maintenance time is up.  It ignores all agents that have less seconds used than min_mature_time,  then it selects num_brains_removed_per_cyle brains to kill off (if possible), then fills up the slots with new brains (the previous brains are killed at the very end, they are just selected to be removed, we still use them for mating, though we do NOT use immature brains)."""
        while self.maintenance_thread_enabled:
            time.sleep(self.genetic_maintenance_time)
            self.brain_lock.acquire()
            mature_agents=set(filter(lambda x: self.seconds_used[x] > self.min_mature_time,range(self.num_brains)))
            if len(mature_agents)==0:
                self.brain_lock.release()
                continue
            #pdb.set_trace()
            self.compute_fitness_list(mature_agents) # This is important to set as other functions access it!
            top_fitness=map(lambda x: [self.true_fitness[x],x],list(mature_agents))
            top_fitness.sort(lambda x,y: int(y[0]-x[0]))
            upper_elite=int(min(self.num_elite_agents,len(top_fitness)+1))
            top_fitness=top_fitness[0:upper_elite]
            top_fitness_agents=set(map(lambda x: x[1],top_fitness)) # Result is a set of elite brain indeces
            # Unset the death_mark of any agent now set elite (man he got lucky I remembered that)
            for i in top_fitness_agents:
                self.death_mark[i]=False
            # See if there are enough mature agents for a good 'ol fashion fight to the death!
            if len(mature_agents) >= self.min_mature_agents:
                # Find how many brains we can kill off
                agents_marked_dead=set(filter(lambda x: self.death_mark[x],range(self.num_brains)))
                max_casualties=self.max_brains_removed_per_cycle-len(agents_marked_dead)
                agents_to_battle=mature_agents-agents_marked_dead-top_fitness_agents
                self.battle_to_the_death(agents_to_battle,max_casualties)

            # Get an updated set of the brains marked dead, then actually kill any of them that aren't checked out 
            killable_brains=set(filter(lambda x: self.death_mark[x] and (not self.checked_out[x]),range(self.num_brains)))
            #pdb.set_trace()
            print('KillableBrains:%s')%(str(killable_brains))
            for i in killable_brains:
                self.initialize_brain_aux(i) # Reset all values to their defaults, except for the family numbers and a new random brain! 
                if random.random() < self.new_member_prob: # Make a totally new member
                    print('B:%d Started a new family') %(i)
                    self.brains[i]=self.__make_new_brain()
                    self.initialize_cur_family(i)
                else:
                    # TODO Consider using roullette wheel
                    # TODO Exclude family members that have too many members from dominant parent selection
                    # Select mature members, do a tournament between them, use this as our dominant parent, then
                    # Select mature members, do a tournament between them, use this as our recessive parent
                    parents=[None]*2
                    temp_mature_list=list(mature_agents)
                    for k in range(2):
                        parents[k]=self.__reducing_tournament_selection(temp_mature_list,self.select_fit_prob)
                        del temp_mature_list[k]
                    print('B:%d Parents(D):%d (R):%d') % (i,parents[0],parents[1])
                    # Initialize family number
                    self.family[i]=self.family[parents[0]]
                    # Perform crossover between states between the two parents
                    temp_brain=self.__crossover_state_agents(parents[0],parents[1])
                    # Mutate the resulting brain
                    self.brains[i]=temp_brain
                    self.__mutate_brain(i)
            self.brain_lock.release()

    def battle_to_the_death(self,brains_to_battle,max_casualties):
        brains_to_battle=set(brains_to_battle)
        num_casualties=0
        #pdb.set_trace()
        while len(brains_to_battle) >=2 and num_casualties < max_casualties:
            contestants=[0]*2
            for i in range(2):
                contestants[i]=random.choice(list(brains_to_battle))
                brains_to_battle-=set([contestants[i]])
            winning_brain=self.__tournament_selection(contestants[0],contestants[1],self.kill_weak_prob)
            losing_brain=contestants[0]
            if winning_brain == contestants[0]:
                losing_brain=contestants[1]
            self.death_mark[losing_brain]=True # Oopsyouredead.com : P
            print('D.M. Agent:%d')%(losing_brain)
            brains_to_battle|=set([winning_brain]) # Get back in there and fight!
            num_casualties+=1
        
    def __make_new_brain(self):
        # Setup weight structure for each state, indexed by layer (0 is connected to the input), node in layer+1, and finally node in layer.
        brain=[]
        for state in range(len(self.state_in_hidden_out)):  # For each state
            weights=[None]*(len(self.state_in_hidden_out[state])-1) # One for each set of weights (e.g., in a 1 hidden layer network, this is len. 2)
            for node_to_layer in range(len(self.state_in_hidden_out[state])-1): # For each 'to' layer (where signal propagates to)
                weights[node_to_layer]=[None]*self.state_in_hidden_out[state][node_to_layer+1]
                for node_to in range(self.state_in_hidden_out[state][node_to_layer+1]): # For each node in the 'to' layer 
                    # We now create the 'from' nodes which includes all nodes in the previous layer plus 1 for the bias
                    weights[node_to_layer][node_to]=[None]*(self.state_in_hidden_out[state][node_to_layer]+1) 
                    for node_from in range(self.state_in_hidden_out[state][node_to_layer]+1): # For each node in the 'from' layer
                        weights[node_to_layer][node_to][node_from]=random.gauss(0,self.initial_weight_std)
            brain.append(weights)
        return brain
        
    def __tournament_selection(self,brain0,brain1,fitter_win_prob):
        """Takes in 2 indeces of brains, and performs tournament selection between them with the fitter winning with a probability of fitter_win_prob."""
        fitter=brain0
        weaker=brain1
        if self.fitness[brain0] < self.fitness[brain1]:
            fitter=brain1
            weaker=brain0
        if random.random() < fitter_win_prob:
            return fitter
        return weaker

    def __reducing_tournament_selection(self,participants,stronger_win_prob):
        part=copy.copy(list(participants))
        random.shuffle(part)
        return reduce(lambda x,y:self.__tournament_selection(x,y,stronger_win_prob),part)

    def __crossover_state_agents(self,brain0, brain1):
        """Takes in 2 indeces of brains, and performs multipoint crossover between them at state boundaries"""
        num_crossed=0
        temp_brain=copy.deepcopy(self.brains[brain0])
        for state in range(len(self.state_in_hidden_out)):
            if random.random() < self.per_state_crossover_prob:
                num_crossed+=1
                temp_brain[state]=copy.deepcopy(self.brains[brain1][state])
        print('CrossedStates: %d') %(num_crossed)
        return temp_brain
    def __mutate_brain(self,brain):
        """Go through every leaf node of the brain and with some probability perturb the weights by some standard deviation"""
        nodes_mutated=0
        for state in range(len(self.state_in_hidden_out)):  # For each state
            for node_to_layer in range(len(self.state_in_hidden_out[state])-1): # For each 'to' layer (where signal propagates to)
                for node_to in range(self.state_in_hidden_out[state][node_to_layer+1]): # For each node in the 'to' layer 
                    for node_from in range(self.state_in_hidden_out[state][node_to_layer]+1): # For each node in the 'from' layer
                        if random.random() < self.per_weight_mutation_prob:
                            nodes_mutated+=1
                            self.brains[brain][state][node_to_layer][node_to][node_from]+=random.gauss(0,self.mutation_weight_std)
        print('NodesMutated: %d') %(nodes_mutated)
